#!/usr/bin/env python3
"""
LangChain智能助手主程序

这是一个功能丰富的LangChain学习项目的主入口。
提供命令行交互界面，展示LangChain的各种功能。
"""

import sys
import os
import asyncio
import json
from typing import Optional, Dict, List, Any
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.prompt import Prompt, Confirm
from rich.table import Table
from rich.live import Live
from rich.spinner import Spinner
from langchain_core.messages import HumanMessage, SystemMessage
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from config.settings import settings
from core.llm_manager import llm_manager
from core.memory_manager import memory_manager
from core.tools_manager import tool_manager
from chains.agent_chain import smart_agent

# 创建Rich控制台
console = Console()

class LangChainAssistant:
    """
    LangChain智能助手主类
    
    提供命令行交互界面和各种功能演示。
    """
    
    def __init__(self):
        """初始化助手"""
        self.console = console
        self.running = True
        self.current_session = "default"
        
        # 系统提示词
        self.system_prompt = """你是一个友好且知识渊博的AI助手，基于LangChain框架构建。
你的任务是帮助用户学习和理解LangChain的各种功能。

请遵循以下原则：
1. 用中文回答问题
2. 提供准确、有用的信息
3. 当不确定时，诚实地说明
4. 保持友好和专业的语调
5. 适当时提供代码示例和解释

你可以帮助用户：
- 回答关于LangChain的问题
- 解释各种概念和功能
- 提供代码示例
- 进行一般性对话
"""
    
    def display_welcome(self):
        """显示欢迎信息"""
        welcome_text = Text()
        welcome_text.append("🦜🔗 ", style="bold blue")
        welcome_text.append("LangChain智能助手", style="bold green")
        welcome_text.append(" v1.0.0", style="dim")
        
        welcome_panel = Panel(
            welcome_text,
            title="欢迎",
            border_style="blue",
            padding=(1, 2)
        )
        
        self.console.print(welcome_panel)
        self.console.print()
        
        # 显示配置信息
        self.display_config_info()
        
        # 显示帮助信息
        self.display_help()
    
    def display_config_info(self):
        """显示配置信息"""
        self.console.print("\n[bold cyan]⚙️ 系统配置信息[/bold cyan]\n")
        
        # LLM 配置
        llm_info = llm_manager.get_model_info()
        llm_table = Table(title="LLM 配置", show_header=True, header_style="bold magenta")
        llm_table.add_column("配置项", style="cyan", no_wrap=True)
        llm_table.add_column("值", style="green")
        
        llm_table.add_row("当前模型", llm_info.get("model_name", "未设置"))
        llm_table.add_row("模型状态", "✅ 就绪" if llm_info.get("is_ready", False) else "❌ 未就绪")
        llm_table.add_row("温度", str(llm_info.get("temperature", "未设置")))
        llm_table.add_row("最大 Tokens", str(llm_info.get("max_tokens", "未设置")))
        
        self.console.print(llm_table)
        
        # 工具配置
        tools_info = tool_manager.get_tool_info()
        tool_config_table = Table(title="工具配置", show_header=True, header_style="bold magenta")
        tool_config_table.add_column("配置项", style="cyan", no_wrap=True)
        tool_config_table.add_column("值", style="green")
        
        tool_config_table.add_row("可用工具数量", str(tools_info["total_tools"]))
        tool_config_table.add_row("工具状态", "✅ 正常" if tools_info["total_tools"] > 0 else "⚠️ 无工具")
        
        self.console.print(tool_config_table)
        
        # 智能体配置
        agent_ready = smart_agent.is_ready()
        agent_table = Table(title="智能体配置", show_header=True, header_style="bold magenta")
        agent_table.add_column("配置项", style="cyan", no_wrap=True)
        agent_table.add_column("值", style="green")
        
        agent_table.add_row("智能体状态", "✅ 就绪" if agent_ready else "❌ 未就绪")
        agent_table.add_row("可用工具", ", ".join(smart_agent.get_available_tools()) if agent_ready else "无")
        
        self.console.print(agent_table)
        
        # 系统信息
        system_table = Table(title="系统信息", show_header=True, header_style="bold magenta")
        system_table.add_column("配置项", style="cyan", no_wrap=True)
        system_table.add_column("值", style="green")
        
        system_table.add_row("Python 版本", f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}")
        system_table.add_row("工作目录", os.getcwd())
        system_table.add_row("配置文件状态", "✅ 已加载" if hasattr(settings, 'config') else "❌ 未加载")
        
        self.console.print(system_table)
        self.console.print()
    
    def display_help(self):
        """显示帮助信息"""
        help_text = """
[bold cyan]可用命令：[/bold cyan]
  [green]/help[/green]     - 显示此帮助信息
  [green]/models[/green]   - 列出可用模型
  [green]/switch[/green]   - 切换模型
  [green]/memory[/green]   - 查看记忆统计
  [green]/clear[/green]    - 清空当前会话
  [green]/history[/green]  - 查看对话历史
  [green]/export[/green]   - 导出会话数据
  [green]/config[/green]   - 显示配置信息
  [green]/tools[/green]    - 查看可用工具
  [green]/agent[/green]    - 切换到智能体模式
  [green]/rag[/green]      - 进入RAG文档问答模式
  [green]/examples[/green] - 运行示例程序
  [green]/quit[/green]     - 退出程序

[bold cyan]使用说明：[/bold cyan]
- 直接输入消息与AI对话
- 使用命令（以/开头）执行特定功能
- 支持多轮对话和上下文记忆
- 智能体模式下可以自动使用工具
- RAG模式下可以基于文档进行问答
- 按Ctrl+C可随时退出
        """
        
        help_panel = Panel(
            help_text,
            title="使用帮助",
            border_style="yellow",
            padding=(1, 2)
        )
        
        self.console.print(help_panel)
        self.console.print()
    
    def handle_command(self, command: str) -> bool:
        """
        处理用户命令
        
        Args:
            command (str): 用户输入的命令
            
        Returns:
            bool: 是否继续运行
        """
        command = command.lower().strip()
        
        if command == "/help":
            self.display_help()
        
        elif command == "/models":
            self.show_available_models()
        
        elif command == "/switch":
            self.switch_model()
        
        elif command == "/memory":
            self.show_memory_stats()
        
        elif command == "/clear":
            self.clear_session()
        
        elif command == "/history":
            self.show_conversation_history()
        
        elif command == "/export":
            self.export_session()
        
        elif command == "/config":
            self.display_config_info()
        
        elif command == "/tools":
            self.show_available_tools()
        
        elif command == "/agent":
            self.switch_to_agent_mode()
        
        elif command == "/rag":
            self.enter_rag_mode()
        
        elif command == "/examples":
            self.show_examples_menu()
        
        elif command == "/quit":
            return False
        
        else:
            self.console.print(f"[red]未知命令: {command}[/red]")
            self.console.print("输入 [green]/help[/green] 查看可用命令")
        
        return True
    
    def show_available_models(self):
        """显示可用模型"""
        models = llm_manager.list_available_models()
        
        if not models:
            self.console.print("[red]没有可用的模型[/red]")
            return
        
        model_table = Table(title="可用模型", show_header=True, header_style="bold magenta")
        model_table.add_column("模型名称", style="cyan", no_wrap=True)
        model_table.add_column("描述", style="green")
        model_table.add_column("状态", style="yellow")
        
        current_model_info = llm_manager.get_model_info()
        current_model_name = current_model_info.get("name", "")
        
        for name, description in models.items():
            status = "当前" if name == current_model_name else ""
            model_table.add_row(name, description, status)
        
        self.console.print(model_table)
        self.console.print()
    
    def switch_model(self):
        """切换模型"""
        models = llm_manager.list_available_models()
        
        if not models:
            self.console.print("[red]没有可用的模型[/red]")
            return
        
        self.console.print("[cyan]可用模型：[/cyan]")
        for i, (name, description) in enumerate(models.items(), 1):
            self.console.print(f"  {i}. {name} - {description}")
        
        try:
            choice = Prompt.ask(
                "请选择模型编号",
                choices=[str(i) for i in range(1, len(models) + 1)],
                default="1"
            )
            
            model_name = list(models.keys())[int(choice) - 1]
            
            with self.console.status(f"正在切换到 {model_name}..."):
                if llm_manager.set_model(model_name):
                    self.console.print(f"[green]已切换到模型: {model_name}[/green]")
                else:
                    self.console.print(f"[red]切换模型失败[/red]")
        
        except (ValueError, IndexError):
            self.console.print("[red]无效的选择[/red]")
    
    def show_memory_stats(self):
        """显示记忆统计"""
        stats = memory_manager.get_memory_stats()
        
        stats_table = Table(title="记忆统计", show_header=True, header_style="bold magenta")
        stats_table.add_column("统计项", style="cyan", no_wrap=True)
        stats_table.add_column("值", style="green")
        
        stats_table.add_row("总会话数", str(stats["total_sessions"]))
        stats_table.add_row("总消息数", str(stats["total_messages"]))
        stats_table.add_row("默认会话ID", stats["default_session_id"])
        
        # 显示记忆类型分布
        for memory_type, count in stats["memory_type_distribution"].items():
            stats_table.add_row(f"{memory_type}记忆", str(count))
        
        self.console.print(stats_table)
        self.console.print()
    
    def clear_session(self):
        """清空当前会话"""
        if Confirm.ask("确定要清空当前会话的历史记录吗？"):
            if memory_manager.clear_session(self.current_session):
                self.console.print("[green]会话历史已清空[/green]")
            else:
                self.console.print("[red]清空会话失败[/red]")
    
    def show_conversation_history(self):
        """显示对话历史"""
        conversation = memory_manager.get_conversation_string(self.current_session)
        
        if not conversation:
            self.console.print("[yellow]当前会话没有历史记录[/yellow]")
            return
        
        history_panel = Panel(
            conversation,
            title="对话历史",
            border_style="blue",
            padding=(1, 2)
        )
        
        self.console.print(history_panel)
        self.console.print()
    
    def export_session(self):
        """导出会话数据"""
        session_data = memory_manager.export_session(self.current_session)
        
        if not session_data:
            self.console.print("[red]导出会话数据失败[/red]")
            return
        
        filename = f"session_{self.current_session}_{session_data['created_at'][:10]}.json"
        
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(session_data, f, ensure_ascii=False, indent=2)
            
            self.console.print(f"[green]会话数据已导出到: {filename}[/green]")
        
        except Exception as e:
            self.console.print(f"[red]导出失败: {str(e)}[/red]")
    
    def show_available_tools(self):
        """显示可用工具"""
        self.console.print("\n[bold cyan]🛠️ 可用工具列表[/bold cyan]\n")
        
        tools_info = tool_manager.get_tool_info()
        
        if tools_info["total_tools"] == 0:
            self.console.print("[red]没有可用的工具[/red]")
            return
        
        tool_table = Table(title="工具清单", show_header=True, header_style="bold magenta")
        tool_table.add_column("工具名称", style="cyan", no_wrap=True)
        tool_table.add_column("描述", style="green")
        tool_table.add_column("状态", style="yellow")
        
        for name, description in tools_info["tool_descriptions"].items():
            status = "✅ 可用" if tool_manager.is_tool_available(name) else "❌ 不可用"
            tool_table.add_row(name, description[:60] + "..." if len(description) > 60 else description, status)
        
        self.console.print(tool_table)
        self.console.print(f"\n总计: {tools_info['total_tools']} 个工具\n")
        
        # 询问是否测试工具
        if Confirm.ask("是否要测试某个工具？"):
            self.test_individual_tool()
    
    def test_individual_tool(self):
        """测试单个工具"""
        tool_names = tool_manager.get_tool_names()
        
        self.console.print("可用工具:", ", ".join(tool_names))
        tool_name = Prompt.ask("请输入要测试的工具名称")
        
        if tool_name not in tool_names:
            self.console.print(f"[red]工具 '{tool_name}' 不存在[/red]")
            return
        
        test_input = Prompt.ask("请输入测试数据")
        
        self.console.print("\n[dim]正在执行工具...[/dim]")

        # 执行工具
        result = tool_manager.execute_tool(tool_name, test_input)
        
        # 显示执行结果
        result_panel = Panel(
            result,
            title=f"工具执行结果: {tool_name}",
            border_style="green",
            padding=(1, 2)
        )
        self.console.print(result_panel)
    
    def switch_to_agent_mode(self):
        """切换到智能体模式"""
        self.console.print("\n[bold cyan]🤖 切换到智能体模式[/bold cyan]\n")
        
        if not smart_agent.is_ready():
            self.console.print("[red]❌ 智能体未准备就绪[/red]")
            self.console.print("[yellow]请检查您的配置和API密钥[/yellow]")
            return
        
        self.console.print("[green]✅ 智能体已准备就绪[/green]")
        self.console.print("[cyan]在智能体模式下，AI可以自动选择和使用工具来回答问题[/cyan]")
        
        # 显示可用工具
        tools = smart_agent.get_available_tools()
        self.console.print(f"[dim]可用工具: {', '.join(tools)}[/dim]\n")
        
        self.console.print("[yellow]智能体对话模式（输入 '/exit' 返回普通模式）[/yellow]\n")
        
        agent_history = []
        
        while True:
            try:
                user_input = Prompt.ask("[bold blue]用户[/bold blue]").strip()
                
                if not user_input:
                    continue
                
                if user_input.lower() == '/exit':
                    self.console.print("[cyan]已退出智能体模式[/cyan]\n")
                    break
                
                # 使用智能体处理
                self.console.print("\n[dim]🤖 智能体正在思考和执行...[/dim]")
                
                with Live(Spinner("dots", "处理中..."), console=self.console):
                    response = smart_agent.run(user_input, agent_history)
                
                # 显示响应
                response_panel = Panel(
                    response,
                    title="🤖 智能体",
                    border_style="green",
                    padding=(1, 2)
                )
                self.console.print(response_panel)
                
                # 更新历史
                from langchain_core.messages import HumanMessage, AIMessage
                agent_history.append(HumanMessage(content=user_input))
                agent_history.append(AIMessage(content=response))
                
                # 保持历史长度合理
                if len(agent_history) > 10:
                    agent_history = agent_history[-8:]
                
                self.console.print()
                
            except KeyboardInterrupt:
                self.console.print("\n[cyan]已退出智能体模式[/cyan]\n")
                break
            except Exception as e:
                self.console.print(f"\n[red]智能体错误: {str(e)}[/red]\n")
    
    def show_examples_menu(self):
        """显示示例程序菜单"""
        self.console.print("\n[bold cyan]📚 示例程序菜单[/bold cyan]\n")
        
        examples = [
            ("基础聊天示例", "examples/basic_chat.py", "展示LangChain的基本对话功能"),
            ("工具使用示例", "examples/tool_example.py", "展示工具的开发和使用"),
            ("智能体示例", "examples/agent_example.py", "展示智能体的功能和推理过程"),
            ("嵌入模型示例", "examples/embedding_example.py", "展示文本嵌入、相似度计算和语义搜索"),
            ("RAG系统示例", "examples/rag_example.py", "展示检索增强生成的完整功能")
        ]
        
        example_table = Table(title="可用示例", show_header=True, header_style="bold magenta")
        example_table.add_column("序号", style="cyan", no_wrap=True)
        example_table.add_column("示例名称", style="green")
        example_table.add_column("描述", style="white")
        
        for i, (name, _, description) in enumerate(examples, 1):
            example_table.add_row(str(i), name, description)
        
        self.console.print(example_table)
        
        try:
            choice = Prompt.ask(f"\n请选择要运行的示例 (1-{len(examples)})", 
                              choices=[str(i) for i in range(1, len(examples) + 1)])
            choice = int(choice) - 1
            
            if 0 <= choice < len(examples):
                example_name, example_file, _ = examples[choice]
                self.console.print(f"\n[cyan]正在运行: {example_name}[/cyan]")
                self.console.print(f"[dim]文件: {example_file}[/dim]\n")
                
                # 运行示例
                try:
                    import subprocess
                    result = subprocess.run([sys.executable, example_file], 
                                          capture_output=False, 
                                          text=True, 
                                          cwd=os.path.dirname(os.path.abspath(__file__)))
                except Exception as e:
                    self.console.print(f"[red]运行示例失败: {str(e)}[/red]")
                    self.console.print(f"[yellow]请手动运行: python {example_file}[/yellow]")
            
        except ValueError:
            self.console.print("[red]无效的选择[/red]")
        except KeyboardInterrupt:
            self.console.print("\n[cyan]已取消[/cyan]")

    async def chat_with_ai(self, user_input: str):
        """
        与AI进行对话
        
        Args:
            user_input (str): 用户输入
        """
        # 添加用户消息到记忆
        memory_manager.add_user_message(user_input, self.current_session)
        
        # 构建消息列表
        messages = [SystemMessage(content=self.system_prompt)]
        
        # 获取历史消息
        history = memory_manager.get_chat_history(self.current_session)
        messages.extend(history)
        
        try:
            if settings.ENABLE_STREAMING:
                # 流式响应
                response_text = ""
                
                with Live(
                    Panel("正在思考...", title="AI助手", border_style="green"),
                    refresh_per_second=10
                ) as live:
                    
                    for chunk in llm_manager.stream(messages):
                        response_text += chunk
                        
                        # 更新显示
                        live.update(
                            Panel(
                                response_text + "▋",
                                title="AI助手",
                                border_style="green",
                                padding=(1, 2)
                            )
                        )
                    
                    # 最终显示
                    live.update(
                        Panel(
                            response_text,
                            title="AI助手",
                            border_style="green",
                            padding=(1, 2)
                        )
                    )
            
            else:
                # 非流式响应
                with self.console.status("正在生成响应..."):
                    response_text = llm_manager.invoke(messages)
                
                response_panel = Panel(
                    response_text,
                    title="AI助手",
                    border_style="green",
                    padding=(1, 2)
                )
                
                self.console.print(response_panel)
            
            # 添加AI响应到记忆
            memory_manager.add_ai_message(response_text, self.current_session)
        
        except Exception as e:
            self.console.print(f"[red]生成响应时出错: {str(e)}[/red]")
            self.console.print("[yellow]请检查您的API密钥和网络连接[/yellow]")
    
    async def run(self):
        """运行主程序"""
        # 验证配置
        if not settings.validate_configuration():
            self.console.print("[red]配置验证失败，请检查您的设置[/red]")
            return
        
        # 显示欢迎信息
        self.display_welcome()
        
        # 主循环
        while self.running:
            try:
                # 获取用户输入
                user_input = Prompt.ask(
                    "[bold blue]您[/bold blue]",
                    default=""
                ).strip()
                
                if not user_input:
                    continue
                
                # 处理命令
                if user_input.startswith('/'):
                    self.running = self.handle_command(user_input)
                    continue
                
                # 与AI对话
                await self.chat_with_ai(user_input)
                
                self.console.print()  # 添加空行
            
            except KeyboardInterrupt:
                if Confirm.ask("\n确定要退出吗？"):
                    break
                else:
                    self.console.print("继续对话...")
            
            except EOFError:
                break
            
            except Exception as e:
                self.console.print(f"[red]发生错误: {str(e)}[/red]")
        
        # 退出信息
        self.console.print("\n[cyan]感谢使用LangChain智能助手！再见！[/cyan]")

    def enter_rag_mode(self):
        """进入RAG文档问答模式"""
        try:
            from core.rag_manager import RAGManager
            
            self.console.print(Panel(
                "🔍 [bold cyan]RAG文档问答模式[/bold cyan]\n\n"
                "在此模式下，您可以：\n"
                "• 添加文档到知识库\n"
                "• 基于文档内容问答\n"
                "• 搜索相关文档\n"
                "• 管理知识库\n\n"
                "输入 [green]help[/green] 查看RAG模式下的可用命令\n"
                "输入 [green]exit[/green] 返回主模式",
                title="RAG模式",
                border_style="cyan"
            ))
            
            # 初始化RAG管理器
            rag_manager = RAGManager()
            
            # RAG模式主循环
            while True:
                try:
                    user_input = Prompt.ask("RAG", default="").strip()
                    
                    if not user_input:
                        continue
                    
                    if user_input.lower() == "exit":
                        self.console.print("[yellow]已退出RAG模式[/yellow]")
                        break
                    
                    elif user_input.lower() == "help":
                        self.show_rag_help()
                    
                    elif user_input.lower() == "status":
                        self.show_rag_status(rag_manager)
                    
                    elif user_input.lower().startswith("add file "):
                        file_path = user_input[9:].strip()
                        self.add_document_to_rag(rag_manager, file_path)
                    
                    elif user_input.lower().startswith("add dir "):
                        dir_path = user_input[8:].strip()
                        self.add_directory_to_rag(rag_manager, dir_path)
                    
                    elif user_input.lower().startswith("add url "):
                        urls = [url.strip() for url in user_input[8:].split(',')]
                        self.add_urls_to_rag(rag_manager, urls)
                    
                    elif user_input.lower().startswith("search "):
                        query = user_input[7:].strip()
                        self.search_rag_documents(rag_manager, query)
                    
                    elif user_input.lower() == "clear":
                        self.clear_rag_knowledge_base(rag_manager)
                    
                    elif user_input.lower() == "export":
                        self.export_rag_summary(rag_manager)
                    
                    else:
                        # 当作问题处理
                        self.ask_rag_question(rag_manager, user_input)
                
                except KeyboardInterrupt:
                    self.console.print("\n[yellow]已退出RAG模式[/yellow]")
                    break
                except Exception as e:
                    self.console.print(f"[red]RAG模式错误: {e}[/red]")
        
        except ImportError:
            self.console.print("[red]RAG功能需要额外的依赖包，请检查安装[/red]")
        except Exception as e:
            self.console.print(f"[red]进入RAG模式失败: {e}[/red]")
    
    def show_rag_help(self):
        """显示RAG模式帮助"""
        help_text = """
[bold cyan]RAG模式命令：[/bold cyan]
  [green]help[/green]                    - 显示此帮助信息
  [green]status[/green]                  - 显示知识库状态
  [green]add file <路径>[/green]         - 添加单个文件到知识库
  [green]add dir <路径>[/green]          - 添加目录中的文档到知识库
  [green]add url <URL1,URL2>[/green]     - 添加网页内容到知识库
  [green]search <查询>[/green]           - 搜索相关文档
  [green]clear[/green]                   - 清空知识库
  [green]export[/green]                  - 导出知识库摘要
  [green]exit[/green]                    - 退出RAG模式

[bold cyan]使用示例：[/bold cyan]
  add file data/documents/example.pdf
  add dir data/documents/
  add url https://example.com/article
  search 人工智能的应用
  什么是机器学习？

[bold cyan]支持的文件格式：[/bold cyan]
  PDF、Word文档、TXT文件、Markdown文件、CSV文件
        """
        
        help_panel = Panel(
            help_text,
            title="RAG模式帮助",
            border_style="cyan",
            padding=(1, 2)
        )
        
        self.console.print(help_panel)
    
    def show_rag_status(self, rag_manager):
        """显示RAG知识库状态"""
        try:
            info = rag_manager.get_knowledge_base_info()
            
            status_table = Table(title="知识库状态", show_header=True, header_style="bold cyan")
            status_table.add_column("项目", style="cyan", no_wrap=True)
            status_table.add_column("值", style="green")
            
            status_table.add_row("初始化状态", "✅ 已初始化" if info['has_vectorstore'] else "❌ 未初始化")
            status_table.add_row("向量数量", str(info['vector_count']))
            status_table.add_row("文档数量", str(info['statistics']['total_documents']))
            status_table.add_row("文本块数量", str(info['statistics']['total_chunks']))
            status_table.add_row("最后更新", info['statistics']['last_updated'] or "未知")
            
            self.console.print(status_table)
            
            if info['statistics']['sources']:
                sources_table = Table(title="数据源", show_header=True, header_style="bold cyan")
                sources_table.add_column("序号", style="cyan", no_wrap=True)
                sources_table.add_column("来源", style="green")
                
                for i, source in enumerate(info['statistics']['sources'], 1):
                    sources_table.add_row(str(i), source)
                
                self.console.print(sources_table)
            else:
                self.console.print("[yellow]知识库中暂无文档[/yellow]")
        
        except Exception as e:
            self.console.print(f"[red]获取知识库状态失败: {e}[/red]")
    
    def add_document_to_rag(self, rag_manager, file_path: str):
        """添加文档文件到RAG知识库"""
        try:
            if not file_path:
                self.console.print("[red]请提供文件路径[/red]")
                return
            
            with self.console.status(f"正在添加文档: {file_path}..."):
                result = rag_manager.add_documents_from_file(file_path)
            
            if result['success']:
                self.console.print(f"[green]✅ {result['message']}[/green]")
            else:
                self.console.print(f"[red]❌ {result['message']}[/red]")
        
        except Exception as e:
            self.console.print(f"[red]添加文档失败: {e}[/red]")
    
    def add_directory_to_rag(self, rag_manager, dir_path: str):
        """添加目录到RAG知识库"""
        try:
            if not dir_path:
                self.console.print("[red]请提供目录路径[/red]")
                return
            
            with self.console.status(f"正在添加目录中的文档: {dir_path}..."):
                result = rag_manager.add_documents_from_directory(dir_path)
            
            if result['success']:
                self.console.print(f"[green]✅ {result['message']}[/green]")
            else:
                self.console.print(f"[red]❌ {result['message']}[/red]")
        
        except Exception as e:
            self.console.print(f"[red]添加目录失败: {e}[/red]")
    
    def add_urls_to_rag(self, rag_manager, urls: List[str]):
        """添加URL到RAG知识库"""
        try:
            if not urls or not urls[0]:
                self.console.print("[red]请提供有效的URL[/red]")
                return
            
            with self.console.status(f"正在添加网页内容: {', '.join(urls)}..."):
                result = rag_manager.add_documents_from_urls(urls)
            
            if result['success']:
                self.console.print(f"[green]✅ {result['message']}[/green]")
            else:
                self.console.print(f"[red]❌ {result['message']}[/red]")
        
        except Exception as e:
            self.console.print(f"[red]添加URL失败: {e}[/red]")
    
    def search_rag_documents(self, rag_manager, query: str):
        """搜索RAG文档"""
        try:
            if not query:
                self.console.print("[red]请提供搜索查询[/red]")
                return
            
            with self.console.status(f"正在搜索: {query}..."):
                results = rag_manager.search_documents(query)
            
            if not results:
                self.console.print("[yellow]未找到相关文档[/yellow]")
                return
            
            search_table = Table(title=f"搜索结果: {query}", show_header=True, header_style="bold cyan")
            search_table.add_column("排名", style="cyan", no_wrap=True, width=6)
            search_table.add_column("来源", style="green", width=30)
            search_table.add_column("内容预览", style="white", width=50)
            
            for result in results:
                content_preview = result['content'][:100] + "..." if len(result['content']) > 100 else result['content']
                search_table.add_row(
                    str(result['rank']),
                    result['source'],
                    content_preview
                )
            
            self.console.print(search_table)
        
        except Exception as e:
            self.console.print(f"[red]搜索失败: {e}[/red]")
    
    def ask_rag_question(self, rag_manager, question: str):
        """向RAG系统提问"""
        try:
            with self.console.status("正在思考..."):
                result = rag_manager.ask_question(question)
            
            if result['success']:
                # 显示回答
                answer_panel = Panel(
                    result['answer'],
                    title="回答",
                    border_style="green",
                    padding=(1, 2)
                )
                self.console.print(answer_panel)
                
                # 显示参考来源
                if result.get('sources'):
                    sources_table = Table(title="参考来源", show_header=True, header_style="bold cyan")
                    sources_table.add_column("序号", style="cyan", no_wrap=True, width=6)
                    sources_table.add_column("来源", style="green", width=30)
                    sources_table.add_column("内容预览", style="white", width=50)
                    
                    for i, source in enumerate(result['sources'], 1):
                        sources_table.add_row(
                            str(i),
                            source['source'],
                            source.get('content_preview', '')
                        )
                    
                    self.console.print(sources_table)
            else:
                self.console.print(f"[red]❌ {result['answer']}[/red]")
        
        except Exception as e:
            self.console.print(f"[red]问答失败: {e}[/red]")
    
    def clear_rag_knowledge_base(self, rag_manager):
        """清空RAG知识库"""
        try:
            if Confirm.ask("确定要清空整个知识库吗？此操作不可撤销"):
                with self.console.status("正在清空知识库..."):
                    success = rag_manager.clear_knowledge_base()
                
                if success:
                    self.console.print("[green]✅ 知识库已清空[/green]")
                else:
                    self.console.print("[red]❌ 清空知识库失败[/red]")
        
        except Exception as e:
            self.console.print(f"[red]清空知识库失败: {e}[/red]")
    
    def export_rag_summary(self, rag_manager):
        """导出RAG知识库摘要"""
        try:
            summary = rag_manager.export_knowledge_base_summary()
            
            summary_panel = Panel(
                summary,
                title="知识库摘要",
                border_style="blue",
                padding=(1, 2)
            )
            self.console.print(summary_panel)
            
            # 询问是否保存到文件
            if Confirm.ask("是否保存摘要到文件？"):
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"knowledge_base_summary_{timestamp}.txt"
                
                try:
                    with open(filename, 'w', encoding='utf-8') as f:
                        f.write(summary)
                    self.console.print(f"[green]✅ 摘要已保存到: {filename}[/green]")
                except Exception as e:
                    self.console.print(f"[red]保存文件失败: {e}[/red]")
        
        except Exception as e:
            self.console.print(f"[red]导出摘要失败: {e}[/red]")

def main():
    """主函数"""
    try:
        assistant = LangChainAssistant()
        asyncio.run(assistant.run())
    
    except KeyboardInterrupt:
        console.print("\n[cyan]程序已退出[/cyan]")
    
    except Exception as e:
        console.print(f"[red]程序启动失败: {str(e)}[/red]")
        console.print("[yellow]请检查您的配置和依赖安装[/yellow]")

if __name__ == "__main__":
    main() 